/* -*- c -*- */

/*
 *****************************************************************************
 **                            INCLUDES                                     **
 *****************************************************************************
 */
#include "Python.h"
#include "numpy/arrayobject.h"
#include "numpy/ufuncobject.h"

/*
 *****************************************************************************
 **                            BASICS                                       **
 *****************************************************************************
 */

typedef npy_intp intp;

#define INIT_OUTER_LOOP_1       \
    intp dN = *dimensions++;    \
    intp N_;                    \
    intp s0 = *steps++;

#define INIT_OUTER_LOOP_2       \
    INIT_OUTER_LOOP_1           \
    intp s1 = *steps++;

#define INIT_OUTER_LOOP_3       \
    INIT_OUTER_LOOP_2           \
    intp s2 = *steps++;

#define INIT_OUTER_LOOP_4       \
    INIT_OUTER_LOOP_3           \
    intp s3 = *steps++;

#define BEGIN_OUTER_LOOP_3      \
    for (N_ = 0; N_ < dN; N_++, args[0] += s0, args[1] += s1, args[2] += s2) {

#define BEGIN_OUTER_LOOP_4      \
    for (N_ = 0; N_ < dN; N_++, args[0] += s0, args[1] += s1, args[2] += s2, args[3] += s3) {

#define END_OUTER_LOOP  }


/*
 *****************************************************************************
 **                             UFUNC LOOPS                                 **
 *****************************************************************************
 */

char *inner1d_signature = "(i),(i)->()";

/**begin repeat

   #TYPE=LONG,DOUBLE#
   #typ=npy_long, npy_double#
*/

/*
 *  This implements the function
 *        out[n] = sum_i { in1[n, i] * in2[n, i] }.
 */
static void
@TYPE@_inner1d(char **args, intp *dimensions, intp *steps, void *func)
{
    INIT_OUTER_LOOP_3
    intp di = dimensions[0];
    intp i;
    intp is1=steps[0], is2=steps[1];
    BEGIN_OUTER_LOOP_3
        char *ip1=args[0], *ip2=args[1], *op=args[2];
        @typ@ sum = 0;
        for (i = 0; i < di; i++) {
            sum += (*(@typ@ *)ip1) * (*(@typ@ *)ip2);
            ip1 += is1;
            ip2 += is2;
        }
        *(@typ@ *)op = sum;
    END_OUTER_LOOP
}

/**end repeat**/

char *innerwt_signature = "(i),(i),(i)->()";

/**begin repeat

   #TYPE=LONG,DOUBLE#
   #typ=npy_long, npy_double#
*/


/*
 *  This implements the function
 *        out[n] = sum_i { in1[n, i] * in2[n, i] * in3[n, i] }.
 */

static void
@TYPE@_innerwt(char **args, intp *dimensions, intp *steps, void *func)
{
    INIT_OUTER_LOOP_4
    intp di = dimensions[0];
    intp i;
    intp is1=steps[0], is2=steps[1], is3=steps[2];
    BEGIN_OUTER_LOOP_4
        char *ip1=args[0], *ip2=args[1], *ip3=args[2], *op=args[3];
        @typ@ sum = 0;
        for (i = 0; i < di; i++) {
            sum += (*(@typ@ *)ip1) * (*(@typ@ *)ip2) * (*(@typ@ *)ip3);
            ip1 += is1;
            ip2 += is2;
            ip3 += is3;
        }
        *(@typ@ *)op = sum;
    END_OUTER_LOOP
}

/**end repeat**/

char *matrix_multiply_signature = "(m,n),(n,p)->(m,p)";

/**begin repeat

   #TYPE=FLOAT,DOUBLE,LONG#
   #typ=npy_float, npy_double,npy_long#
*/

/*
 *  This implements the function
 *        out[k, m, p] = sum_n { in1[k, m, n] * in2[k, n, p] }.
 */


static void
@TYPE@_matrix_multiply(char **args, intp *dimensions, intp *steps, void *func)
{
    /* no BLAS is available */
    INIT_OUTER_LOOP_3
    intp dm = dimensions[0];
    intp dn = dimensions[1];
    intp dp = dimensions[2];
    intp m,n,p;
    intp is1_m=steps[0], is1_n=steps[1], is2_n=steps[2], is2_p=steps[3],
         os_m=steps[4], os_p=steps[5];
    intp ib1_n = is1_n*dn;
    intp ib2_n = is2_n*dn;
    intp ib2_p = is2_p*dp;
    intp ob_p  = os_p *dp;
    BEGIN_OUTER_LOOP_3
        char *ip1=args[0], *ip2=args[1], *op=args[2];
        for (m = 0; m < dm; m++) {
            for (n = 0; n < dn; n++) {
                register @typ@ val1 = (*(@typ@ *)ip1);
                for (p = 0; p < dp; p++) {
                    if (n == 0) *(@typ@ *)op = 0;
                    *(@typ@ *)op += val1 * (*(@typ@ *)ip2);
                    ip2 += is2_p;
                    op  +=  os_p;
                }
                ip2 -= ib2_p;
                op  -=  ob_p;
                ip1 += is1_n;
                ip2 += is2_n;
            }
            ip1 -= ib1_n;
            ip2 -= ib2_n;
            ip1 += is1_m;
            op  +=  os_m;
        }
    END_OUTER_LOOP
}

/**end repeat**/

/*  The following lines were generated using a slightly modified
    version of code_generators/generate_umath.py and adding these
    lines to defdict:

defdict = {
'inner1d' :
    Ufunc(2, 1, None_,
        r'''inner on the last dimension and broadcast on the rest \n"
        "     \"(i),(i)->()\" \n''',
          TD('ld'),
          ),
'innerwt' :
    Ufunc(3, 1, None_,
          r'''inner1d with a weight argument \n"
          "     \"(i),(i),(i)->()\" \n''',
          TD('ld'),
          ),
}

*/

static PyUFuncGenericFunction inner1d_functions[] = { LONG_inner1d, DOUBLE_inner1d };
static void * inner1d_data[] = { (void *)NULL, (void *)NULL };
static char inner1d_signatures[] = { PyArray_LONG, PyArray_LONG, PyArray_LONG, PyArray_DOUBLE, PyArray_DOUBLE, PyArray_DOUBLE };
static PyUFuncGenericFunction innerwt_functions[] = { LONG_innerwt, DOUBLE_innerwt };
static void * innerwt_data[] = { (void *)NULL, (void *)NULL };
static char innerwt_signatures[] = { PyArray_LONG, PyArray_LONG, PyArray_LONG, PyArray_LONG, PyArray_DOUBLE, PyArray_DOUBLE, PyArray_DOUBLE, PyArray_DOUBLE };
static PyUFuncGenericFunction matrix_multiply_functions[] = { LONG_matrix_multiply, FLOAT_matrix_multiply, DOUBLE_matrix_multiply };
static void *matrix_multiply_data[] = { (void *)NULL, (void *)NULL, (void *)NULL };
static char matrix_multiply_signatures[] = { PyArray_LONG, PyArray_LONG, PyArray_LONG,  PyArray_FLOAT, PyArray_FLOAT, PyArray_FLOAT,  PyArray_DOUBLE, PyArray_DOUBLE, PyArray_DOUBLE };

static void
addUfuncs(PyObject *dictionary) {
    PyObject *f;

    f = PyUFunc_FromFuncAndDataAndSignature(inner1d_functions, inner1d_data, inner1d_signatures, 2,
                                    2, 1, PyUFunc_None, "inner1d",
                                    "inner on the last dimension and broadcast on the rest \n"\
                                    "     \"(i),(i)->()\" \n",
                                    0, inner1d_signature);
    PyDict_SetItemString(dictionary, "inner1d", f);
    Py_DECREF(f);
    f = PyUFunc_FromFuncAndDataAndSignature(innerwt_functions, innerwt_data, innerwt_signatures, 2,
                                    3, 1, PyUFunc_None, "innerwt",
                                    "inner1d with a weight argument \n"\
                                    "     \"(i),(i),(i)->()\" \n",
                                    0, innerwt_signature);
    PyDict_SetItemString(dictionary, "innerwt", f);
    Py_DECREF(f);
    f = PyUFunc_FromFuncAndDataAndSignature(matrix_multiply_functions,
                                    matrix_multiply_data, matrix_multiply_signatures,
                                    3, 2, 1, PyUFunc_None, "matrix_multiply",
                                    "matrix multiplication on last two dimensions \n"\
                                    "     \"(m,n),(n,p)->(m,p)\" \n",
                                    0, matrix_multiply_signature);
    PyDict_SetItemString(dictionary, "matrix_multiply", f);
    Py_DECREF(f);
}

/*
    End of auto-generated code.
*/



static PyObject *
UMath_Tests_test_signature(PyObject *dummy, PyObject *args)
{
    int nin, nout;
    PyObject *signature;
    PyObject *f;
    int core_enabled;

    if (!PyArg_ParseTuple(args, "iiO", &nin, &nout, &signature)) return NULL;
    f = PyUFunc_FromFuncAndDataAndSignature(NULL, NULL, NULL,
        0, nin, nout, PyUFunc_None, "no name",
        "doc:none",
        1, PyString_AS_STRING(signature));
    if (f == NULL) return NULL;
    core_enabled = ((PyUFuncObject*)f)->core_enabled;
    return Py_BuildValue("i", core_enabled);
}

static PyMethodDef UMath_TestsMethods[] = {
    {"test_signature",  UMath_Tests_test_signature, METH_VARARGS,
     "Test signature parsing of ufunc. \n"
     "Arguments: nin nout signature \n"
     "If fails, it returns NULL. Otherwise it will returns 0 for scalar ufunc "
     "and 1 for generalized ufunc. \n",
     },
    {NULL, NULL, 0, NULL}        /* Sentinel */
};

PyMODINIT_FUNC
initumath_tests(void)
{
    PyObject *m;
    PyObject *d;
    PyObject *version;

    m = Py_InitModule("umath_tests", UMath_TestsMethods);
    if (m == NULL) return;

    import_array();
    import_ufunc();

    d = PyModule_GetDict(m);

    version = PyString_FromString("0.1");
    PyDict_SetItemString(d, "__version__", version);
    Py_DECREF(version);

    /* Load the ufunc operators into the module's namespace */
    addUfuncs(d);

    if (PyErr_Occurred()) {
        PyErr_SetString(PyExc_RuntimeError,
                        "cannot load umath_tests module.");
    }
}
